代码实现 RoPE 与 YaRN

从零理解 RoPE 与 YaRN：结合 PyTorch 源码详解位置编码实现

在学习大模型源码时，很多人第一次看到 RoPE（Rotary Position Embedding）和 YaRN（Yet another RoPE extensioN）的实现，都会被各种 cos、sin、freqs、ramp 搞晕。

实际上，这段代码背后的思想并不复杂：

RoPE：让向量按照不同频率进行旋转，从而编码位置信息。
YaRN：在超长上下文场景下，降低部分旋转频率，避免位置混叠（Aliasing）。

下面先看下RoPE和YaRN的实现逻辑：

import math
import torch

def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
    freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
    if rope_scaling is not None: # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
        orig_max, factor, beta_fast, beta_slow, attn_factor = (
            rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
            rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
        )
        if end / orig_max > 1.0:
            inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
            low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
            ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
            freqs = freqs * (1 - ramp + ramp / factor)
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
    return freqs_cos, freqs_sin

def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
    q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
    k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
    return q_embed, k_embed

一、为什么需要位置编码

Transformer 的注意力机制本身不具备顺序概念。

例如：

我 爱 中国
中国 爱 我

对于 Attention 来说：

QK^T

只是一堆向量之间的相似度计算。

模型并不知道：

第一个词
第二个词
第三个词

谁在前，谁在后。

因此必须引入位置编码（Position Embedding）。

二、RoPE 的核心思想

传统位置编码通常是：

x_{i} + p_{i}

即：

词向量 + 位置向量

RoPE 采用完全不同的思路：

不直接添加位置，而是让向量旋转。

二维空间中的旋转

假设有一个二维向量：

x = [\begin{matrix} x_{1} \\ x_{2} \end{matrix}]

位置为 $m$ 时：

R_{m} = [\begin{matrix} \cos (m θ) & - \sin (m θ) \\ \sin (m θ) & \cos (m θ) \end{matrix}]

则：

x_{m} = R_{m} x

图示：

位置0

      ●

位置1

       ●
      /

位置2

      ●
     /
    /

位置越远

旋转角度越大

三、RoPE 为什么有效

Attention 实际计算：

q_{m}^{T} k_{n}

经过旋转后：

(R_{m} q)^{T} (R_{n} k)

利用旋转矩阵性质：

R_{m}^{T} R_{n} = R_{n - m}

得到：

(R_{m} q)^{T} (R_{n} k) = q^{T} R_{n - m} k

注意：

结果只依赖于：

n - m

即：

两个 Token 的相对距离

而不是绝对位置。

这正是 Transformer 所需要的信息。

四、源码解析：构造频率

源码：

freqs = 1.0 / (
    rope_base **
    (
        torch.arange(0, dim, 2)
        .float() / dim
    )
)

假设：

dim = 8
rope_base = 10000

那么：

torch.arange(0, dim, 2)

得到：

[0,2,4,6]

计算后：

维度组      频率

0,1         1
2,3         0.1
4,5         0.01
6,7         0.001

数学形式：

第 $i$ 对维度：

θ_{i} = \frac{1}{b a s e^{2 i / d}}

可以发现：

低维度频率高
高维度频率低

图示：

高频

↻↻↻↻↻↻↻

中频

↻↻↻

低频

↻

超低频

·

五、构造所有位置的旋转角度

源码：

t = torch.arange(end)
freqs = torch.outer(t, freqs)

假设：

位置：

0
1
2
3

频率：

1
0.1
0.01

那么：

      1     0.1    0.01

0     0      0      0

1     1      0.1    0.01

2     2      0.2    0.02

3     3      0.3    0.03

实际上是在计算：

angle = p o s i t i o n \times f r e q u e n c y

即：

m θ_{i}

六、计算 sin 与 cos

源码：

freqs_cos = torch.cos(freqs)

freqs_sin = torch.sin(freqs)

得到：

\cos (m θ_{i})

和

\sin (m θ_{i})

继续：

torch.cat(
    [torch.cos(freqs),
     torch.cos(freqs)],
    dim=-1
)

为什么要复制一次？

因为：

一个频率对应两个维度

例如：

(x1,x2)

共用一个角度。

七、真正的旋转实现

核心代码：

def rotate_half(x):
    return torch.cat(
        (
            -x[..., x.shape[-1]//2:],
            x[..., :x.shape[-1]//2]
        ),
        dim=-1
    )

例如：

x = [a,b,c,d]

变成：

[-c,-d,a,b]

实际上等价于：

(x, y) \to (- y, x)

这正是二维旋转矩阵中的关键部分。

八、应用 RoPE

源码：

q_embed = q * cos + rotate_half(q) * sin

同理：

k_embed = k * cos + rotate_half(k) * sin

写成数学形式：

q^{'} = q \cos θ + q_{⊥} \sin θ

其中：

q_{⊥} = (- q_{2}, q_{1})

展开：

\begin{aligned} q_{1}^{'} & = q_{1} \cos θ - q_{2} \sin θ \\ q_{2}^{'} & = q_{1} \sin θ + q_{2} \cos θ \end{aligned}

这正是：

R (θ) q

因此：

apply_rotary_pos_emb()

本质上就是：

给 Q 和 K 做旋转

九、RoPE 的问题

假设训练长度：

最大角度：

2048 θ_{i}

推理时：

长度。

则：

32768 θ_{i}

高频维度会疯狂旋转：

↻↻↻↻↻↻↻↻↻↻↻↻

出现：

位置混叠（Aliasing）

例如：

位置10000

位置10200

旋转后的结果可能非常接近。

模型开始分不清位置。

十、YaRN 的思想

YaRN 的核心非常简单：

长上下文时，让部分频率变慢。

原频率：

f (i)

变成：

f^{'} (i) = \frac{f (i)}{s}

其中：

s = 16

例如：

训练长度

2048

推理长度

32768

则：

32768 / 16 = 2048

角度重新回到训练区间。

十一、为什么不能全部缩放

如果全部频率都缩放：

f^{'} (i) = f (i) / 16

会导致：

邻近 Token 的差异变小

局部感知能力下降。

因此 YaRN 采用：

低频保持

高频缩放

中间平滑过渡

十二、YaRN 源码解析

核心代码：

ramp = torch.clamp(
	(
		torch.arange(
				dim // 2, device=freqs.device
			).float() - low) 
			/ 
		max(high - low, 0.001), 
		0,
		1
	)

得到：

维度:

0 1 2 3 4 5 6 7

ramp:

0 0 0 .2 .5 .8 1 1

含义：

低维度：

不缩放

高维度：

完全缩放

中间：

线性过渡

然后：

freqs = freqs *
	(
		 1 - ramp
		 + ramp / factor
	)

数学形式：

f^{'} (i) = f (i) ((1 - γ_{i}) + \frac{γ_{i}}{s})

其中：

γ_{i} \in [0, 1]

即：

ramp

当：

γ = 0

时：

f^{'} = f

当：

γ = 1

时：

f^{'} = \frac{f}{s}

图示：

原频率

│\
│ \
│  \
│   \
└──────

YaRN

│\
│ \
│  \____
│
└────────

高频部分被压缩。

十三、beta_fast 与 beta_slow 的作用

源码：

low  = inv_dim(beta_fast)
high = inv_dim(beta_slow)

用于确定：

从哪一维开始缩放

到哪一维完成缩放

最终形成：

低频区域
↓↓↓↓↓↓↓

保持原样

██████████

过渡区域

▒▒▒▒▒▒▒▒▒▒

高频区域

缩放16倍

□□□□□□□□□

十四、整体流程总结

第一步：生成频率

freqs

得到：

θ_{i} = \frac{1}{b a s e^{2 i / d}}

第二步：计算角度

position × frequency

得到：

m θ_{i}

第三步：计算

sin
cos

第四步：旋转 Q、K

q * cos + rotate(q) * sin

得到：

q^{'} = R (θ) q

k^{'} = R (θ) k

第五步：YaRN 扩展

对高频维度进行缩放：

f^{'} (i) = f (i) ((1 - γ_{i}) + \frac{γ_{i}}{s})

从而避免长上下文下：

角度爆炸

和

位置混叠

结语

从代码层面看：

RoPE 本质上是一组不同频率的旋转器（Rotators）：

a n g l e = p o s i t i o n \times f r e q u e n c y

然后利用：

\cos (a n g l e)

和

\sin (a n g l e)

把位置信息编码进 Q、K。

而 YaRN 本质上是一种频率重映射（Frequency Rescaling）策略：

f (i) \to f^{'} (i)

通过压缩高频旋转速度，把原本只能处理 2K 上下文的模型扩展到 32K、64K 甚至更长，同时尽量保持短距离建模能力。

从信号处理视角看：

RoPE = 多频率相位编码（Multi-frequency Phase Encoding）
YaRN = 高频压缩（High-frequency Compression）

两者结合，构成了当前主流大模型（Llama、Qwen、DeepSeek 等）长上下文能力的重要基础。